#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import division

import matplotlib.pyplot
import os
import numpy
import sys
import ConfigParser
import math
import cPickle as pickle
import exact_diagonalization as ED

config = ConfigParser.RawConfigParser()
config.read('config.py') #opening configuration file
modelname = 'Triangular_transverse_2d'
ferromagnetic = config.getboolean(modelname, 'ferromagnetic')

if ferromagnetic:
    case = 'ferromagnetic'
    J = 1
else:
    case = 'antiferromagnetic'
    J = -1

if '-B' in sys.argv:
    B = sys.argv[ sys.argv.index('-B') + 1 ]
    
N = int(sys.argv[ sys.argv.index('-N') + 1 ])

tRange = numpy.arange(0.1, 10, 0.1) 

fName = os.path.join('accurate_data', modelname, case, 'N{0}_J{1}_B{2}'.format(N, J, float(B)) )
fName = open(fName)

#Energies = numpy.loadtxt(fName)
data = pickle.load(fName)
Energies = data[0]
minE = min(data[0])
eigvec = data[1]

nStates = (N**2)**2 

tE = []
for T in tRange:
    Z = 0
    E = 0
    E2 = 0
    for Energy in Energies:
        mnog = math.exp( -( Energy - minE )/ T )
        #mnog = math.exp( - Energy/ T )
        Z += mnog
        E += Energy * mnog
        E2 += Energy**2 * mnog
        
    aE = E / Z
    aE2 = E2 / Z
    if '-C' in sys.argv:
        
        tE += [ (aE2 - aE**2) / T**2/ N**2]
    elif '-E' in sys.argv:
        tE += [ aE / N**2 ]
    elif '-M' in sys.argv:
        M = {}
        temp = 0
        for n in xrange( nStates ):
            M[n] = 0
            for j, V in enumerate(eigvec):
                M[n] += V[n]**2 * ( sum( ED.i2state(j, N**2) ) / N**2 )**2
            temp += M[n] * math.exp( - ( Energies[n] - minE) / T)
        tE += [temp / Z]
         
        
if '-C' in sys.argv:
    fName = open('C_N{0}_B{1}'.format(N, B), 'w')

elif '-E' in sys.argv:
    fName = open('E_N{0}_B{1}'.format(N, B), 'w')
elif '-M' in sys.argv:
    fName = open('M_N{0}_B{1}'.format(N, B), 'w')

print >> fName, (list(tRange), list(tE))
fName.close()
    
matplotlib.pyplot.plot(tRange, tE)
matplotlib.pyplot.show()    